# -*- coding: utf-8 -*-  
'''
gpt2.0相关预处理

@author: luoyi
Created on 2021年4月7日
'''
import tensorflow as tf
import pickle
import numpy as np

import utils.conf as conf


#    位置掩码
def position_mask(sentence_maxlen):
    '''
        @param sentence_maxlen: 句子最大长度
        @return [1, 1, sentence_maxlen, sentence_maxlen]的反下三角阵
                    0的地方不需要掩，1的地方需要掩
    '''
    #    [sentence_maxlen * sentence_maxlen的下三角阵]
    mask = tf.ones(shape=(sentence_maxlen, sentence_maxlen), dtype=tf.float32)
    mask = 1 - tf.linalg.band_part(mask, -1, 0)
    #    扩展为[batch_size, n_heads, sentence_maxlen, sentence_maxlen]的格式，用的时候自行广播
    mask = mask[tf.newaxis, tf.newaxis, :, :]
    return mask


#    填充掩码
def padding_mask(inputs):
    '''
        @param inputs: 输入数据 Tensor(batch_size, sentence_maxlen)
                            sentence_maxlen维中非0值是实际的词编码，0值为填充<PAD>
        @return: [batch_size, 1, 1, sentence_maxlen]
                    0的地方不需要掩，1的地方需要掩
    '''
    mask = tf.cast(tf.equal(inputs, 0), dtype=tf.float32)         #    Tensor(batch_size, sentence_maxlen)
    #    扩展为[batch_size, 1, 1, sentence_maxlen]，用的时候自行广播
    mask = mask[:, tf.newaxis, tf.newaxis, :]
    return mask


#    loss掩码
def loss_mask(y_true):
    '''
        @param y_true: Tensor(batch_size, sentence_maxlen)
        @return: 掩码矩阵mask: Tensor(batch_size, sentence_maxlen)
                                有词的地方为1，<PAD>为0
                 每个batch_size实际词数量: Tensor(batch_size, )
    '''
    mask = tf.cast(tf.not_equal(y_true, 0), dtype=tf.float32)
    effective_count = tf.math.count_nonzero(mask, axis=-1)
    return mask, effective_count


#    加载预训练的词向量参数
def pre_embedding_weights(fpath=conf.GPT2.get_embedding_weights_path()):
    try:
        with open(file=fpath, mode='rb') as fr:
            embedding_weights = pickle.load(fr)
            embedding_weights = np.array(embedding_weights)
            embedding_weights = np.squeeze(embedding_weights)
            pass
        return embedding_weights
    except:
        print('读取词向量pkl文件异常. fpath:', fpath)
        return None


